import tkinter as tk
from openai import OpenAI


class ChatApp:
    def __init__(self, root):
        self.root = root
        self.root.title("Chat with OpenAI")

        self.message_frame = tk.Frame(root)
        self.message_frame.pack(padx=10, pady=10)

        self.messages_text = tk.Text(self.message_frame, wrap="word", state="disabled")
        self.messages_text.pack(expand=True, fill="both")

        self.input_frame = tk.Frame(root)
        self.input_frame.pack(padx=10, pady=(0, 10))

        self.input_entry = tk.Entry(self.input_frame)
        self.input_entry.pack(side="left", expand=True, fill="x")
        self.input_entry.bind("<Return>", lambda event: self.send_message())

        self.send_button = tk.Button(self.input_frame, text="发送", command=self.send_message)
        self.send_button.pack(side="right")

        self.client = OpenAI(
            base_url="https://integrate.api.nvidia.com/v1",
            api_key="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"
        )

        self.messages_all = []

    def send_message(self):
        user_input = self.input_entry.get().strip()
        if user_input:
            self.messages_all.append({
                "content": user_input,
                "role": "user"
            })
            completion = self.client.chat.completions.create(
                model="google/codegemma-7b",
                messages=self.messages_all,
                temperature=0.5,
                top_p=1,
                max_tokens=1024,
                stream=True
            )
            response_msg = ""
            self.messages_text.config(state="normal")
            self.messages_text.insert("end", f"You: {user_input}\n", "user")
            for chunk in completion:
                if chunk.choices[0].delta.content is not None:
                    temp_msg = chunk.choices[0].delta.content
                    response_msg += temp_msg
            self.messages_all.append({
                "content": response_msg,
                "role": "assistant"
            })
            self.messages_text.insert("end", f"Assistant: {response_msg}\n", "assistant")
            self.messages_text.config(state="disabled")
            self.messages_text.see("end")
            self.input_entry.delete(0, "end")


root = tk.Tk()
app = ChatApp(root)
root.mainloop()
# 我希望user的信息居右，和assistant区分开